#!/usr/bin/env python 
# -*- coding:utf-8 -*-

from ensemble_boxes import *
import json
import os
from tqdm import tqdm
import numpy as np

iou_thr = 0.7
weights = None

save_path = "results/front_ensemble.json"
label_file = "data/annotations/front_test.json"

labels = json.load(open(label_file, 'r'))

imgid2size = {}
info = []
bboxs_model = [{}, {}, {}]
score_model = [{}, {}, {}]
class_model = [{}, {}, {}]

for img in labels["images"]:
    imgid2size[img["id"]] = (img["height"], img["width"])

anno_file1 = "results/front/model1.bbox.json"
anno_file2 = "results/front/model2.bbox.json"
anno_file3 = "results/front/model3.bbox.json"


anno_list = [json.load(open(anno_file1, 'r'))
    , json.load(open(anno_file2, 'r'))
    , json.load(open(anno_file3, 'r'))]

num_models = len(anno_list)

for i, annos in enumerate(anno_list):
    for anno in annos:
        if anno["image_id"] not in bboxs_model[i].keys():
            bboxs_model[i][anno["image_id"]] = []
            score_model[i][anno["image_id"]] = []
            class_model[i][anno["image_id"]] = []

        x, y, w, h = anno["bbox"]
        height, width = imgid2size[anno["image_id"]]
        normbox = [x / width, y / height, (x + w) / width, (y + h) / height]
        bboxs_model[i][anno["image_id"]].append(normbox)
        score_model[i][anno["image_id"]].append(anno["score"])
        class_model[i][anno["image_id"]].append(anno["category_id"])

for id in tqdm(bboxs_model[0]):
    bboxs_list = []
    scores_list = []
    class_list = []
    for i in range(num_models):
        bboxs_model[i][id], score_model[i][id], class_model[i][id] = nms([bboxs_model[i][id]], [score_model[i][id]],
                                                                         [class_model[i][id]], weights=None,
                                                                         iou_thr=0.6)
        bboxs_list.append(bboxs_model[i][id])
        scores_list.append(score_model[i][id])
        class_list.append(class_model[i][id])

    boxes, scores, classes = weighted_boxes_fusion(bboxs_list, scores_list, class_list, weights=weights,
                                                   iou_thr=iou_thr)
    for i in range(boxes.shape[0]):
        temp = {}
        x1, y1, x2, y2 = boxes[i]
        hr, wr = imgid2size[id]
        realbox = [x1 * wr, y1 * hr, (x2 - x1) * wr, (y2 - y1) * hr]

        # print('id:', id)
        # print('boxes[i]:', boxes[i])
        # print('wr,hr:', wr, hr)
        # print('realbox:', realbox)
        # print('score:', scores[i])
        # print('category_id:', int(classes[i]))
        # print('=============================================================================')

        temp["image_id"] = id
        temp["bbox"] = realbox
        temp["score"] = float(scores[i])
        temp["category_id"] = int(classes[i])
        info.append(temp)


json_file = json.dumps(info)
with open(save_path, "w", encoding='utf-8') as f:
    f.write(json_file)